// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/dwio/common/on_demand_unit_loader.h>

#include <numeric>

#include <pollux/common/base/exceptions.h>
#include <pollux/dwio/common/measure_time.h>
#include <pollux/dwio/common/unit_loader_tools.h>

using kumo::pollux::dwio::common::measureTimeIfCallback;

namespace kumo::pollux::dwio::common {
    namespace {
        class OnDemandUnitLoader : public UnitLoader {
        public:
            OnDemandUnitLoader(
                std::vector<std::unique_ptr<LoadUnit> > loadUnits,
                std::function<void(std::chrono::high_resolution_clock::duration)>
                blockedOnIoCallback)
                : loadUnits_{std::move(loadUnits)},
                  blockedOnIoCallback_{std::move(blockedOnIoCallback)} {
            }

            ~OnDemandUnitLoader() override = default;

            LoadUnit &getLoadedUnit(uint32_t unit) override {
                POLLUX_CHECK_LT(unit, loadUnits_.size(), "Unit out of range");

                if (loadedUnit_.has_value()) {
                    if (loadedUnit_.value() == unit) {
                        return *loadUnits_[unit];
                    }

                    loadUnits_[*loadedUnit_]->unload();
                    loadedUnit_.reset();
                } {
                    auto measure = measureTimeIfCallback(blockedOnIoCallback_);
                    loadUnits_[unit]->load();
                }
                loadedUnit_ = unit;

                return *loadUnits_[unit];
            }

            void onRead(uint32_t unit, uint64_t rowOffsetInUnit, uint64_t /* rowCount */)
            override {
                POLLUX_CHECK_LT(unit, loadUnits_.size(), "Unit out of range");
                POLLUX_CHECK_LT(
                    rowOffsetInUnit, loadUnits_[unit]->getNumRows(), "Row out of range");
            }

            void onSeek(uint32_t unit, uint64_t rowOffsetInUnit) override {
                POLLUX_CHECK_LT(unit, loadUnits_.size(), "Unit out of range");
                POLLUX_CHECK_LE(
                    rowOffsetInUnit, loadUnits_[unit]->getNumRows(), "Row out of range");
            }

        private:
            const std::vector<std::unique_ptr<LoadUnit> > loadUnits_;
            const std::function<void(std::chrono::high_resolution_clock::duration)>
            blockedOnIoCallback_;
            std::optional<uint32_t> loadedUnit_;
        };
    } // namespace

    std::unique_ptr<UnitLoader> OnDemandUnitLoaderFactory::create(
        std::vector<std::unique_ptr<LoadUnit> > loadUnits,
        uint64_t rowsToSkip) {
        const auto totalRows = std::accumulate(
            loadUnits.cbegin(), loadUnits.cend(), 0UL, [](uint64_t sum, auto &unit) {
                return sum + unit->getNumRows();
            });
        POLLUX_CHECK_LE(
            rowsToSkip,
            totalRows,
            "Can only skip up to the past-the-end row of the file.");
        return std::make_unique<OnDemandUnitLoader>(
            std::move(loadUnits), blockedOnIoCallback_);
    }
} // namespace kumo::pollux::dwio::common
